import sqlite3
import networkx as nx
from typing import Dict, List, Any, Tuple
from pathlib import Path
from ..core.logger import get_logger

logger = get_logger('graph_rag.financial_kg')

class FinancialKnowledgeGraph:
    """金融知识图谱构建器"""
    
    def __init__(self, db_path: str = None):
        """
        初始化金融知识图谱
        
        Args:
            db_path: 数据库文件路径
        """
        self.graph = nx.DiGraph()
        
        # 默认数据库路径
        if db_path is None:
            self.db_path = Path(__file__).parent.parent.parent / "dataset" / "sqlite" / "博金杯比赛数据.db"
        else:
            self.db_path = Path(db_path)
        
        if not self.db_path.exists():
            logger.error(f"数据库文件不存在: {self.db_path}")
            raise FileNotFoundError(f"数据库文件不存在: {self.db_path}")
        
        logger.info(f"Financial knowledge graph initialized with database: {self.db_path}")
    
    def build_knowledge_graph(self):
        """构建金融知识图谱"""
        try:
            # 连接数据库
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row
            
            # 构建各种实体关系
            self._build_fund_entities(conn)
            self._build_stock_entities(conn)
            self._build_bond_entities(conn)
            self._build_relationships(conn)
            
            conn.close()
            logger.info(f"知识图谱构建完成，包含 {self.graph.number_of_nodes()} 个节点和 {self.graph.number_of_edges()} 条边")
            
        except Exception as e:
            logger.error(f"构建知识图谱时出错: {e}")
            raise
    
    def _build_fund_entities(self, conn: sqlite3.Connection):
        """构建基金实体"""
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT 基金代码, 基金简称, 基金类型, 管理人, 托管人 FROM 基金基本信息 LIMIT 1000")
            funds = cursor.fetchall()
            
            for fund in funds:
                fund_code = fund['基金代码']
                fund_name = fund['基金简称']
                
                # 添加基金节点
                self.graph.add_node(
                    f"fund_{fund_code}",
                    type="fund",
                    name=fund_name,
                    code=fund_code,
                    fund_type=fund['基金类型'],
                    manager=fund['管理人'],
                    custodian=fund['托管人']
                )
                
        except Exception as e:
            logger.error(f"构建基金实体时出错: {e}")
    
    def _build_stock_entities(self, conn: sqlite3.Connection):
        """构建股票实体"""
        try:
            # 获取A股信息
            cursor = conn.cursor()
            cursor.execute("SELECT DISTINCT 股票代码 FROM A股票日行情表 LIMIT 1000")
            stocks = cursor.fetchall()
            
            for stock in stocks:
                stock_code = stock['股票代码']
                
                # 添加A股节点
                self.graph.add_node(
                    f"stock_{stock_code}",
                    type="stock",
                    code=stock_code,
                    market="A股"
                )
            
            # 获取港股信息
            cursor.execute("SELECT DISTINCT 股票代码 FROM 港股票日行情表 LIMIT 1000")
            hk_stocks = cursor.fetchall()
            
            for stock in hk_stocks:
                stock_code = stock['股票代码']
                
                # 添加港股节点
                self.graph.add_node(
                    f"hk_stock_{stock_code}",
                    type="stock",
                    code=stock_code,
                    market="港股"
                )
                
        except Exception as e:
            logger.error(f"构建股票实体时出错: {e}")
    
    def _build_bond_entities(self, conn: sqlite3.Connection):
        """构建债券实体"""
        try:
            # 获取基金债券持仓信息
            cursor = conn.cursor()
            cursor.execute("""
                SELECT DISTINCT 债券名称, 债券类型 
                FROM 基金债券持仓明细 
                WHERE 债券名称 IS NOT NULL AND 债券类型 IS NOT NULL
                LIMIT 1000
            """)
            bonds = cursor.fetchall()
            
            for bond in bonds:
                bond_name = bond['债券名称']
                bond_type = bond['债券类型']
                
                # 添加债券节点
                self.graph.add_node(
                    f"bond_{bond_name}",
                    type="bond",
                    name=bond_name,
                    bond_type=bond_type
                )
                
        except Exception as e:
            logger.error(f"构建债券实体时出错: {e}")
    
    def _build_relationships(self, conn: sqlite3.Connection):
        """构建实体间关系"""
        try:
            # 构建基金-股票持仓关系
            self._build_fund_stock_relationships(conn)
            
            # 构建基金-债券持仓关系
            self._build_fund_bond_relationships(conn)
            
            # 构建基金-可转债持仓关系
            self._build_fund_convertible_bond_relationships(conn)
            
        except Exception as e:
            logger.error(f"构建实体关系时出错: {e}")
    
    def _build_fund_stock_relationships(self, conn: sqlite3.Connection):
        """构建基金-股票持仓关系"""
        try:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT 基金代码, 股票代码, 市值, 市值占基金资产净值比 
                FROM 基金股票持仓明细 
                WHERE 基金代码 IS NOT NULL AND 股票代码 IS NOT NULL
                LIMIT 2000
            """)
            holdings = cursor.fetchall()
            
            for holding in holdings:
                fund_code = holding['基金代码']
                stock_code = holding['股票代码']
                
                # 检查基金和股票节点是否存在
                fund_node = f"fund_{fund_code}"
                stock_node = f"stock_{stock_code}"
                
                if fund_node in self.graph.nodes and stock_node in self.graph.nodes:
                    # 添加持仓关系
                    self.graph.add_edge(
                        fund_node,
                        stock_node,
                        relation="holds",
                        market_value=holding['市值'],
                        ratio=holding['市值占基金资产净值比']
                    )
                    
        except Exception as e:
            logger.error(f"构建基金-股票关系时出错: {e}")
    
    def _build_fund_bond_relationships(self, conn: sqlite3.Connection):
        """构建基金-债券持仓关系"""
        try:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT 基金代码, 债券名称, 持债市值, 持债市值占基金资产净值比 
                FROM 基金债券持仓明细 
                WHERE 基金代码 IS NOT NULL AND 债券名称 IS NOT NULL
                LIMIT 2000
            """)
            holdings = cursor.fetchall()
            
            for holding in holdings:
                fund_code = holding['基金代码']
                bond_name = holding['债券名称']
                
                # 检查基金和债券节点是否存在
                fund_node = f"fund_{fund_code}"
                bond_node = f"bond_{bond_name}"
                
                if fund_node in self.graph.nodes and bond_node in self.graph.nodes:
                    # 添加持仓关系
                    self.graph.add_edge(
                        fund_node,
                        bond_node,
                        relation="holds",
                        market_value=holding['持债市值'],
                        ratio=holding['持债市值占基金资产净值比']
                    )
                    
        except Exception as e:
            logger.error(f"构建基金-债券关系时出错: {e}")
    
    def _build_fund_convertible_bond_relationships(self, conn: sqlite3.Connection):
        """构建基金-可转债持仓关系"""
        try:
            cursor = conn.cursor()
            cursor.execute("""
                SELECT 基金代码, 债券名称, 市值, 市值占基金资产净值比 
                FROM 基金可转债持仓明细 
                WHERE 基金代码 IS NOT NULL AND 债券名称 IS NOT NULL
                LIMIT 1000
            """)
            holdings = cursor.fetchall()
            
            for holding in holdings:
                fund_code = holding['基金代码']
                bond_name = holding['债券名称']
                
                # 检查基金和债券节点是否存在
                fund_node = f"fund_{fund_code}"
                bond_node = f"bond_{bond_name}"
                
                if fund_node in self.graph.nodes and bond_node in self.graph.nodes:
                    # 添加持仓关系
                    self.graph.add_edge(
                        fund_node,
                        bond_node,
                        relation="holds_convertible_bond",
                        market_value=holding['市值'],
                        ratio=holding['市值占基金资产净值比']
                    )
                    
        except Exception as e:
            logger.error(f"构建基金-可转债关系时出错: {e}")
    
    def query_entity_relationships(self, entity_id: str, max_depth: int = 2) -> Dict:
        """
        查询实体的关系信息
        
        Args:
            entity_id: 实体ID
            max_depth: 最大搜索深度
            
        Returns:
            关系信息字典
        """
        if entity_id not in self.graph.nodes:
            return {"error": f"实体 {entity_id} 不存在"}
        
        # 获取实体基本信息
        entity_info = dict(self.graph.nodes[entity_id])
        
        # 获取直接关联的实体
        neighbors = list(self.graph.neighbors(entity_id))
        relationships = []
        
        for neighbor in neighbors:
            edge_data = self.graph.get_edge_data(entity_id, neighbor)
            neighbor_info = dict(self.graph.nodes[neighbor])
            
            relationships.append({
                "target": neighbor_info,
                "relationship": edge_data
            })
        
        return {
            "entity": entity_info,
            "relationships": relationships
        }
    
    def find_paths_between_entities(self, source_id: str, target_id: str, max_paths: int = 5) -> List[List]:
        """
        查找两个实体之间的路径
        
        Args:
            source_id: 源实体ID
            target_id: 目标实体ID
            max_paths: 最大路径数
            
        Returns:
            路径列表
        """
        try:
            # 使用networkx查找所有简单路径
            paths = list(nx.all_simple_paths(self.graph, source_id, target_id, cutoff=max_paths))
            return paths
        except nx.NetworkXNoPath:
            return []
        except Exception as e:
            logger.error(f"查找路径时出错: {e}")
            return []